import json
import os
import re
from typing import Any, List


def extract_code(input: str) -> str | None:
    # Use regular expressions to find Java code blocks
    pattern = r"```(java)?(.*?)```"
    matches = re.findall(pattern, input, re.DOTALL)
    for match in matches:
        if "@Test" in match[1]:
            return match[1].strip()
    return None


def extract_json(input: str, expected_schema=None) -> list[Any] | None | Any:
    # Used to store all valid JSON objects
    json_objects = []

    # 1. If the entire string is in JSON format, parse it directly
    try:
        # Try to parse the entire string directly
        parsed_json = json.loads(input)
        if expected_schema and not validate_json_structure(parsed_json, expected_schema):
            return None, True
        json_objects.append(parsed_json)
    except json.JSONDecodeError:
        pass  # If the entire string cannot be parsed as JSON, try other ways

    # 2. Extract parts wrapped in triple backticks and try to parse each part
    # Regex matches ` ```json ... ``` ` or ` ``` ... ``` `
    code_block_pattern = r'```(json)?(.*?)```'  # Matches triple backtick content with or without json identifier
    code_blocks = re.findall(code_block_pattern, input, flags=re.DOTALL)  # `re.DOTALL` allows matching newlines
    find_json = False

    for block in code_blocks:
        block_content = block[1].strip()  # Get the content inside backticks and remove leading/trailing whitespace
        try:
            # Try to parse the content of each code block
            parsed_json = json.loads(block_content)
            find_json = True
            if expected_schema and not validate_json_structure(parsed_json, expected_schema):
                continue  # If the format or type does not match expectations, skip this JSON
            json_objects.append(parsed_json)
        except json.JSONDecodeError:
            pass  # If it cannot be parsed as JSON, ignore this code block

    # If there is valid JSON data, return the list of JSON objects, otherwise return None
    return json_objects if json_objects else None, find_json


def validate_json_structure(json_obj, expected_schema):
    """
    Validate whether the structure and type of the JSON object meet expectations.

    :param json_obj: The JSON object (dictionary) to validate.
    :param expected_schema: The expected JSON structure, in the format {key: expected_type}, for example:
                            {"name": str, "value": int, "data": list}
    :return: Returns True if all keys' types match expectations; otherwise returns False.
    """
    # Iterate through each key in the expected structure
    if json_obj is None:
        return False
    for key, expected_type in expected_schema.items():
        # Check if the key exists and the type matches
        if key not in json_obj:
            print(f"Missing key: {key}")
            return False

        if not isinstance(json_obj[key], expected_type):
            print(f"Incorrect type for key '{key}': Expected {expected_type}, got {type(json_obj[key])}")
            return False

    return True

def add_record(log_dir, name, ret, find_bugs, coverage, mutation_score, line_coverage_rates, branch_coverage_rates,
               mutation_scores, total_lines=set(), covered_lines=set(), mutation_info={}, update=False):
    with open(os.path.join(log_dir, "record.jsonl"), "a") as f:
        f.write(json.dumps(
            {"update": update, "name": name, "ret": ret, "find_bugs": find_bugs, "coverage": coverage, "mutation_score": mutation_score,
             "line_coverage_rates": line_coverage_rates, "branch_coverage_rates": branch_coverage_rates,
             "mutation_scores": mutation_scores, "total_lines": list(total_lines), "covered_lines": list(covered_lines),
             "mutation_info": mutation_info}) + "\n")


def add_test_case(log_dir, name, find_bugs, test_point, test_case, signature, mutants_list=[]):
    with open(os.path.join(log_dir, "test_case.jsonl"), "a") as f:
        f.write(json.dumps({"signature": signature, "name": name, "find_bugs": find_bugs, "test_point": test_point,
                             "test_case": test_case, "mutants_list": mutants_list}) + "\n")

if __name__ == '__main__':
    # Test the extract_json function
    input_text = """
    ```json
    {
        "name": "Alice",
        "age": 25
    }
    ```
    
    ```
    {
        "name": "Bob",
        "age": 31
    }
    ```
    
    ```python
    print("Hello, World!")
    ```
    """
    json_objects = extract_json(input_text)
    if json_objects:
        for obj in json_objects:
            print(obj)
    else:
        print("No valid JSON objects found.")
